'''
-*- coding: utf-8 -*-
@File  : waterpipe_service.py
@Author: Shanmh
@Time  : 2024/02/29 下午6:11
@Function：
'''
import os.path
import time

import cv2
from ultralytics import YOLO
from datetime import datetime
from loguru import logger
from  tools.Rtsp import PushStream,LoadStreams,LoadImages
from tools.detection import Detection

import io
import contextlib

def algo(model,frames,ori_img):
    results = model.predict(frames, device="cpu")
    # cv2.imshow("a",frames[0])
    # cv2.waitKey(1)
    for ind_,item in enumerate(results):
        image = item.orig_img
        boxes = item.boxes.cpu().numpy()  # get boxes on cpu in numpy
        for ind, box in enumerate(boxes):  # iterate boxes
            r = box.xyxy[0].astype(int)
            cls = boxes.cls[ind]  #
            cv2.putText(image, f"{cls}", r[:2], color=(155, 15, 0), fontFace=1, fontScale=5, thickness=3)
            cv2.rectangle(image, r[:2], r[2:], (0, 0, 255), 2)  # draw boxes on img

        # ori_img[ind_][area_box[0][0]:area_box[1][0], area_box[0][1]:area_box[1][1], :]=image[:,:,:].copy()
        # ori_img = cv2.resize(ori_img, (900, 600))
        # cv2.imshow(f"show0", ori_img)
        # cv2.waitKey(1)

        if len(results[ind_])>0 and True:
            pass
            # imgname = saveImg("water")
            # # cv2.imwrite(f"{os.path.join(RootPath().rootpath, imgname)}", frames[ind])
            # cv2.imwrite(f"{os.path.join(RootPath().rootpath, imgname)}",cv2.resize(ori_img[ind_], (640, 480)))
            #

area_box=[]
logger.add(sink='log/log_{time}_post.log',filter=lambda x: '[post]' in x['message'], enqueue=True,rotation='00:00', retention="7 days")
logger.add(sink='log/log_{time}_pull.log',filter=lambda x: '[pull]' in x['message'], enqueue=True,rotation='00:00', retention="7 days")
logger.add(sink='log/log_{time}_log.log',filter=lambda x: '[fps]' in x['message'], enqueue=True,rotation='00:00', retention="7 days")


config={"video_url":"rtsp://admin:hxzh2019@192.168.6.135:554/h264/ch1/main/av_stream"}
def main():
    global config
    model = YOLO("models/yolov8s.pt")
    if config["video_url"].startswith("rtsp"):
        load_stream=LoadStreams(config["video_url"],buffer=True)
    else:
        load_stream= LoadImages(config["video_url"])
    last_time=0
    fps=0

    det_images=[]
    ori_images=[]

    for i in load_stream:
        (rtsp_url, images, _, _) = i[:4]
        det_img=images[0]
        fps+=1
        if time.time()-last_time>1:
            logger.info(f"[fps] ： {fps}")
            last_time=time.time()
            fps=0
        # imgname = saveImg("water")
        # cv2.imwrite(f"{os.path.join(RootPath().rootpath, imgname)}", det_img)
        # continue

        if len(det_images)>=20:
            algo(model,det_images,ori_images)
            det_images=[]
            ori_images=[]
        else:
            det_images.append(det_img)
            ori_images.append(images[0])




if __name__ == '__main__':
    main()
    pass

